/*****************************************************************************
 * http_proxy_DNS.go
 * Names:
 * NetIds:
 *****************************************************************************/

// TODO: implement an HTTP proxy with DNS Prefetching

// Note: it is highly recommended to complete http_proxy.go first, then copy it
// with the name http_proxy_DNS.go, thus overwriting this file, then edit it
// to add DNS prefetching (don't forget to change the filename in the header
// to http_proxy_DNS.go in the copy of http_proxy.go)

package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"

	"golang.org/x/net/html"
)

// 解析 HTML，提取出所有 <a href="..."> 标签中的链接
func parseHTMLLinks(body []byte) []string {
	var links []string
	tokenizer := html.NewTokenizer(strings.NewReader(string(body)))

	for {
		tt := tokenizer.Next()
		switch tt {
		case html.ErrorToken:
			return links
		case html.StartTagToken, html.SelfClosingTagToken:
			t := tokenizer.Token()
			if t.Data == "a" {
				for _, attr := range t.Attr {
					if attr.Key == "href" && strings.HasPrefix(attr.Val, "http") {
						links = append(links, attr.Val)
					}
				}
			}
		}
	}
}

// 执行 DNS 查询（仅发起查询，不处理响应）
func dnsQuery(url string, wg *sync.WaitGroup) {
	defer wg.Done()

	// 发起 DNS 查询
	parsedURL, err := net.ResolveIPAddr("ip", url)
	if err != nil {
		fmt.Println("Error resolving DNS for", url, ":", err)
		return
	}

	// 这里只需要发起 DNS 查询，忽略返回值
	fmt.Printf("DNS query sent for %s, resolved to %s\n", url, parsedURL.String())
}

// 处理 HTTP 请求并并发执行 DNS 查询
func handleRequest(clientConn net.Conn) {
	defer clientConn.Close()

	// 读取 HTTP 请求
	buffer := make([]byte, 4096)
	n, err := clientConn.Read(buffer)
	if err != nil {
		fmt.Println("Error reading from client:", err)
		return
	}
	request := string(buffer[:n])

	// 解析请求的 URI
	requestMethod := ""
	requestURI := ""
	if idx := strings.Index(request, "GET"); idx != -1 {
		requestMethod = "GET"
		requestURI = strings.TrimSpace(strings.Split(request, " ")[1])
	} else {
		io.WriteString(clientConn, "500 Internal Error")
		return
	}

	// 处理 GET 请求
	if requestMethod == "GET" {
		// 发起 HTTP 请求
		// req, err := http.NewRequest("GET", requestURI, nil)

		req, err := http.NewRequest("GET", requestURI, nil)
		if err != nil {
			fmt.Println("Error creating request:", err)
			return
		}
		// req.Header.Set("Host", requestURI)
		// req.Header.Set("Connection", "close")
		// // resp, err := http.DefaultClient.Do(req)
		// host := req.Host
		// targetConn, err := net.Dial("tcp", host+":80")

		// if err != nil {
		// 	io.WriteString(clientConn, "500 Internal Error")
		// 	return
		// }
		// defer resp.Body.Close()
		host := strings.Split(requestURI, "/")[2] // 获取域名部分
		targetConn, err := net.Dial("tcp", host+":80")
		if err != nil {
			io.WriteString(clientConn, "500 Internal Error")
			return
		}

		defer targetConn.Close()

		req.Header.Set("Connection", "close")
		err = req.Write(targetConn)
		if err != nil {
			io.WriteString(clientConn, fmt.Sprintf("500 Internal Error"))
			return
		}
		// httpRequest := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nConnection: close\r\n\r\n", requestURI, host)
		// _, err = targetConn.Write([]byte(httpRequest))
		// if err != nil {
		// 	io.WriteString(clientConn, "500 Internal Error")
		// 	return
		// }

		responseBuffer := new(bytes.Buffer)
		// teeReader := io.TeeReader(targetConn, responseBuffer)

		io.Copy(clientConn, targetConn)

		// 读取完整响应
		bodyBytes, err := ioutil.ReadAll(responseBuffer)
		if err != nil {
			fmt.Println("Error reading response body:", err)
			return
		}

		// 解析 HTML 获取链接
		links := parseHTMLLinks(bodyBytes)

		// 使用 WaitGroup 并发执行 DNS 查询
		var wg sync.WaitGroup
		for _, link := range links {
			wg.Add(1)
			go dnsQuery(link, &wg) // 在新的 goroutine 中执行 DNS 查询
		}
		wg.Wait() // 等待所有 DNS 查询完成

		// err = req.Write(targetConn)
		// if err != nil {
		// 	io.WriteString(clientConn, fmt.Sprintf("500 Internal Error"))
		// 	return
		// }
		// io.Copy(clientConn, teeReader)

		// // 在新 goroutine 中解析 HTML 和执行 DNS 查询
		// bodyBytes, err := ioutil.ReadAll(resp.Body) // Go 1.16+ 使用 io.ReadAll
		// if err != nil {
		// 	fmt.Println("Error reading response body:", err)
		// 	return
		// }

		// // 解析 HTML 获取链接
		// links := parseHTMLLinks(bodyBytes)

		// // 使用 WaitGroup 并发执行 DNS 查询
		// var wg sync.WaitGroup
		// for _, link := range links {
		// 	wg.Add(1)
		// 	go dnsQuery(link, &wg) // 在新的 goroutine 中执行 DNS 查询
		// }
		// wg.Wait() // 等待所有 DNS 查询完成
	}
}

func main() {
	// 监听指定端口
	if len(os.Args) < 2 {
		fmt.Println("Please provide the port to listen on")
		return
	}
	port := ":" + os.Args[1]

	// 开始监听端口
	listener, err := net.Listen("tcp", port)
	if err != nil {
		fmt.Println("Error listening:", err)
		return
	}
	defer listener.Close()
	fmt.Println("Listening on", port)

	// 接受并处理连接
	for {
		clientConn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			continue
		}
		fmt.Println("Connection accepted")

		// 使用 goroutine 处理请求
		go handleRequest(clientConn)
	}
}
